from fastapi import APIRouter, Body
from typing import Dict, Any
from intelli_port.commons.observability import ab_exposure_total
from intelli_port.commons.observability import counter_sum_by_label
from intelli_port.commons.observability import Counter
from prometheus_client import Counter as PcCounter


router = APIRouter()


# success counter per experiment/bucket
ab_success_total: Counter = PcCounter(
    "ab_success_total",
    "AB experiment successes",
    ["experiment_id", "bucket"],
)


@router.post("/assign")
async def assign(payload: Dict[str, Any] = Body(...)):
    eid = str(payload.get("experiment_id") or "").strip()
    buckets = payload.get("buckets") or []
    if not eid or not buckets:
        return {"bucket": None, "alloc": []}
    import random
    alloc = []
    best_b = None
    best_score = -1.0
    for b in buckets:
        expo = counter_sum_by_label(ab_exposure_total, "experiment_id", eid)
        expo_b = counter_sum_by_label(ab_exposure_total, "bucket", str(b))
        succ_b = counter_sum_by_label(ab_success_total, "bucket", str(b))
        succ_e_b = succ_b
        trials = max(1.0, expo_b)
        alpha = 1.0 + float(succ_e_b)
        beta = 1.0 + max(0.0, float(trials - succ_e_b))
        score = random.betavariate(alpha, beta)
        alloc.append({"bucket": str(b), "score": score})
        if score > best_score:
            best_score = score
            best_b = str(b)
    return {"bucket": best_b, "alloc": alloc}


@router.post("/mark")
async def mark(payload: Dict[str, Any] = Body(...)):
    eid = str(payload.get("experiment_id") or "").strip()
    bucket = str(payload.get("bucket") or "").strip()
    success = bool(payload.get("success") or False)
    if not eid or not bucket:
        return {"ok": False}
    ab_exposure_total.labels(experiment_id=eid, bucket=bucket).inc()
    if success:
        ab_success_total.labels(experiment_id=eid, bucket=bucket).inc()
    return {"ok": True}


@router.get("/metrics")
async def metrics(eid: str):
    expos = counter_sum_by_label(ab_exposure_total, "experiment_id", eid)
    return {"experiment_id": eid, "exposures": expos}